import os
os.environ["OPENAI_API_KEY"] = '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'
from langchain import PromptTemplate, FewShotPromptTemplate
from langchain.llms import OpenAI

examples = [
    {"word": "黑", "antonym": "白"},
    {"word": "伤心", "antonym": "开心"},
]

example_template = """
单词: {word}
反义词: {antonym}\\n
"""

# 创建提示词模版
example_prompt = PromptTemplate(
    input_variables=["word", "antonym"],
    template=example_template,
)

# 创建小样本提示词模版
few_shot_prompt = FewShotPromptTemplate(
    examples=examples,
    example_prompt=example_prompt,
    prefix="给出每个单词的反义词",
    suffix="单词: {input}\\n反义词:",
    input_variables=["input"],
    example_separator="\\n",
)

# 格式化小样本提示词
prompt_text = few_shot_prompt.format(input="粗")

# 调用OpenAI
llm = OpenAI(temperature=0.9)

print(llm(prompt_text))